'use strict';

const { Document } = require('/document.js');
const { DocumentCommand, CompoundCommandBuilder, AddChildNodesCommandBuilder, NodeChildType, NodeMoveType } = require('/commands.js');
const { Dialog, DialogResult } = require('/dialog.js');
const { UnitType } = require('/units.js');
const { FillDescriptor } = require('/fills.js');
const { RGBA8 } = require('/colours.js');
const { Selection } = require('/selections.js');
const { ContainerNodeDefinition } = require('/nodes.js');
const { app } = require('/application.js');

function abort(msg) { app.alert(msg); throw new Error(msg); }
function isOk(r) { return (r?.value ?? r) === DialogResult.Ok.value; }

function toHex(r, g, b) {
    return '#' + [r, g, b].map(v => v.toString(16).padStart(2, '0')).join('').toUpperCase();
}

const doc = Document.current;
if (!doc) abort('No active document.');
const spread = doc.currentSpread;

// Recursively collect all solid-fill vector nodes
function collectSolidFillNodes(parent, result, depth) {
    if (depth > 10 || !parent.children) return;
    for (let i = 0; i < parent.children.length; i++) {
        const n = parent.children.at(i);
        const fd = n.brushFillInterface?.fillDescriptor;
        const fill = fd?.fill;
        if (fill && fill.constructor.name === 'SolidFill') {
            const rgba = fill.colour?.rgba8;
            if (rgba) result.push({ node: n, r: rgba.r, g: rgba.g, b: rgba.b });
        } else {
            collectSolidFillNodes(n, result, depth + 1);
        }
    }
}

const selNodes = doc.selection.nodes;
const shapeNodes = [];
if (selNodes.length > 0) {
    for (let i = 0; i < selNodes.length; i++) {
        const n = selNodes.at(i);
        const fd = n.brushFillInterface?.fillDescriptor;
        const fill = fd?.fill;
        if (fill && fill.constructor.name === 'SolidFill') {
            const rgba = fill.colour?.rgba8;
            if (rgba) shapeNodes.push({ node: n, r: rgba.r, g: rgba.g, b: rgba.b });
        }
        collectSolidFillNodes(n, shapeNodes, 0);
    }
} else {
    collectSolidFillNodes(spread, shapeNodes, 0);
}

if (shapeNodes.length === 0) abort('No solid-fill vector shapes found.\nSelect a group or layer containing solid-filled vector shapes.');
console.log(`Found ${shapeNodes.length} solid-fill nodes`);

// Default palette
const defaults = [
    RGBA8(255,   0,   0, 255),
    RGBA8(  0, 180,   0, 255),
    RGBA8(  0,   0, 255, 255),
    RGBA8(255, 255,   0, 255),
    RGBA8(255, 128,   0, 255),
    RGBA8(128,   0, 200, 255),
    RGBA8(200, 200, 200, 255),
];

// ---- Dialog ----
const dlg = Dialog.create('Colour Reduce Vectors');
const col = dlg.addColumn();

const gN = col.addGroup('Colour Swap');
const nC = gN.addUnitValueEditor('Number of colours (1–7)', UnitType.Number, UnitType.Number, 7, 1, 7).setPrecision(0);
const groupCheck = gN.addCheckBox('Group curves by colour', true);

const pk = [];
for (let i = 0; i < 7; i++) {
    pk.push(gN.addColourPicker(`Colour ${i + 1}`, defaults[i]));
}

dlg.initialWidth = 320;
const res = dlg.runModal();
if (!isOk(res)) abort('Cancelled.');

const numColours = Math.max(1, Math.min(7, Math.round(nC.value)));
const doGrouping = groupCheck.value;

// Read active palette
const palette = [];
for (let i = 0; i < numColours; i++) {
    const rgba = pk[i].value.rgba8;
    palette.push({ r: rgba.r, g: rgba.g, b: rgba.b });
}
console.log('Palette:', palette.map(c => toHex(c.r, c.g, c.b)).join(' | '));

doc.executeCommand(DocumentCommand.createSetCurrentSpread(spread));

function colourDist(a, b) {
    const dr=a.r-b.r, dg=a.g-b.g, db=a.b-b.b;
    return dr*dr + dg*dg + db*db;
}

// Assign each node to its nearest palette colour
const clusters = Array.from({ length: numColours }, () => []);
for (const sn of shapeNodes) {
    let best=0, bestDist=Infinity;
    for (let i=0; i<palette.length; i++) {
        const d = colourDist(sn, palette[i]);
        if (d<bestDist) { bestDist=d; best=i; }
    }
    clusters[best].push(sn.node);
}

// Step 1: Create colour group containers if requested
let containers = null;
if (doGrouping) {
    doc.executeCommand(DocumentCommand.createSetCurrentSpread(spread));
    const gb = AddChildNodesCommandBuilder.create();
    gb.setInsertionTarget(spread);
    for (let i=0; i<numColours; i++) {
        if (clusters[i].length > 0) {
            const p = palette[i];
            gb.addContainerNode(ContainerNodeDefinition.create(`Colour ${i + 1} ${toHex(p.r, p.g, p.b)}`));
        }
    }
    const gc = gb.createCommand(false, NodeChildType.Main);
    doc.executeCommand(gc);

    // Map containers back to palette indices
    containers = new Array(numColours).fill(null);
    let ci = 0;
    const created = [];
    gc.newNodes.forEach(n => created.push(n));
    for (let i=0; i<numColours; i++) {
        if (clusters[i].length > 0) containers[i] = created[ci++];
    }
    console.log(`Created ${created.length} colour group containers`);
}

// Step 2: Recolour + move into groups — all in ONE compound = single undo step
doc.executeCommand(DocumentCommand.createSetCurrentSpread(spread));
const compound = CompoundCommandBuilder.create();

for (let i=0; i<numColours; i++) {
    const p = palette[i];
    const fill = FillDescriptor.createSolid(RGBA8(p.r, p.g, p.b, 255));
    for (const node of clusters[i]) {
        compound.addCommand(DocumentCommand.createSetBrushFill(
            Selection.create(doc, node), fill
        ));
        if (doGrouping && containers?.[i]) {
            compound.addCommand(DocumentCommand.createMoveNodes(
                Selection.create(doc, node), containers[i], NodeMoveType.Inside, NodeChildType.Main
            ));
        }
    }
}

doc.executeCommand(compound.createCommand());
console.log(`Done! Recoloured ${shapeNodes.length} shapes to ${numColours} colours — single undo step.`);